package com.sunzm.flink.datastream.java.source;

import org.apache.commons.lang3.RandomUtils;
import org.apache.commons.lang3.time.DateFormatUtils;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.apache.kafka.clients.consumer.ConsumerConfig;

import java.sql.*;
import java.util.ArrayList;
import java.util.Properties;

/**
 * FLink的各种source示例代码
 */
public class JavaFlinkSourceDemo {
    private static boolean isLocal = true;

    public static void main(String[] args) throws Exception {
        //获取执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        if(isLocal){
            env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());
        }

        //从一个集合创建 DataStream
        //fromCollection(env);

        //直接把若干个元素创建成 DataStream
        //fromElements(env);

        //读取一个文件，创建DataStream
        //readTextFile(env);

        //从一个socket端口实时接收数据
        //socketTextStream(env);

        //自定义数据源
        //cusSourceSimple(env);

        //自定义数据源,从MySQL获取数据
        //cusSourceMySQL(env);

        //kafka数据源
        Properties properties = new Properties();
        properties.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "82.156.210.70:9093");
        properties.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "flink-test");
        properties.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest");
        DataStream<String> dataStream = env.addSource(new FlinkKafkaConsumer<>("my-topic", new SimpleStringSchema(), properties));

        dataStream.print();

        //启动程序
        String className = Thread.currentThread().getStackTrace()[1].getClassName();
        String simpleName = className.substring(className.lastIndexOf(".") + 1);
        env.execute(simpleName);
    }

    /**
     * 集合类型的数据源
     *
     * @param env
     */
    private static void fromCollection(StreamExecutionEnvironment env) {
        ArrayList<String> list = new ArrayList<>();
        list.add("hello,java,hello,hadoop");
        list.add("hello, spark,hello, flink");

        //从数据源获取数据
        DataStream<String> dataStream = env.fromCollection(list);

        //将结果打印输出
        dataStream.print();
    }

    /**
     * 元素类型的数据源
     *
     * @param env
     */
    private static void fromElements(StreamExecutionEnvironment env) {

        //从数据源获取数据
        DataStream<Integer> dataStream = env.fromElements(1, 3, 5, 7, 9);

        //将结果打印输出
        dataStream.print();
    }

    /**
     * socket端口数据源
     *
     * @param env
     */
    private static void socketTextStream(StreamExecutionEnvironment env) {

        //从数据源获取数据
        DataStream<String> dataStream = env.socketTextStream("82.156.210.70", 9999);

        //将结果打印输出
        dataStream.print();
    }

    /**
     * 文本文件数据源
     *
     * @param env
     */
    private static void readTextFile(StreamExecutionEnvironment env) {

        //从数据源获取数据
        DataStream<String> dataStream = env.readTextFile("data/spark/rdd/word.txt", "UTF-8");

        //将结果打印输出
        dataStream.print();
    }

    /**
     * 用户自定义数据源
     */
    private static void cusSourceSimple(StreamExecutionEnvironment env) {
        //从数据源获取数据
        DataStream<String> dataStream = env.addSource(new MyDataSourceSimple());

        //将结果打印输出
        dataStream.print();
    }

    /**
     * 用户自定义数据源，从MySQl获取数据
     */
    private static void cusSourceMySQL(StreamExecutionEnvironment env) {
        //从数据源获取数据
        DataStream<String> dataStream = env.addSource(new MyDataSourceMySQL());

        //将结果打印输出
        dataStream.print();
    }

    private static class MyDataSourceSimple implements SourceFunction<String> {
        private boolean isRunning = true;

        /**
         * 程序启动时，会调用一次这个方法
         *
         * @param ctx
         * @throws Exception
         */
        @Override
        public void run(SourceContext<String> ctx) throws Exception {
            while (isRunning) {
                //获取当前时间
                long currentTimeMillis = System.currentTimeMillis();

                String formatDateStr = DateFormatUtils.format(currentTimeMillis, "yyyy-MM-dd HH:mm:ss");

                //将数据输出
                ctx.collect(formatDateStr);

                //防止数据产生过快，休眠一会
                int nextInt = RandomUtils.nextInt(100, 20000);
                Thread.sleep(nextInt);
            }
        }

        /**
         * 程序取消时会调用一次这个方法
         */
        @Override
        public void cancel() {
            isRunning = false;
        }
    }


    private static class MyDataSourceMySQL extends RichSourceFunction<String> {
        private boolean isRunning = true;
        private Connection connection = null;
        private PreparedStatement pstmt = null;
        private ResultSet rs = null;
        private String sql = "SELECT * FROM t_user";

        /**
         * 程序初始化时，会运行一次,我们可以在这里创建数据库连接，这样就只会创建一次
         *
         * @param parameters
         * @throws Exception
         */
        @Override
        public void open(Configuration parameters) throws Exception {

            connection = DriverManager.getConnection("jdbc:mysql://localhost:3306/mydb?useSSL=false", "root", "root");

            pstmt = connection.prepareStatement(sql);
        }

        @Override
        public void run(SourceContext<String> ctx) throws Exception {
            while (isRunning) {
                rs = pstmt.executeQuery();
                while (rs.next()) {
                    String userId = rs.getString("userId");
                    String name = rs.getString("name");
                    String age = rs.getString("age");

                    //输出结果
                    ctx.collect(userId + "," + name + "," + age);
                }

                //防止数据频繁访问数据源，休息10秒
                Thread.sleep(10000);
            }
        }

        /**
         * 程序取消时会调用一次这个方法
         */
        @Override
        public void cancel() {
            closeConn();
        }

        /**
         * 程序停止时，会运行一次
         *
         * @throws Exception
         */
        @Override
        public void close() throws Exception {
            closeConn();
        }

        private void closeConn() {
            isRunning = false;
            if (rs != null) {
                try {
                    rs.close();
                } catch (SQLException throwables) {
                    throwables.printStackTrace();
                }
            }
            if (pstmt != null) {
                try {
                    pstmt.close();
                } catch (SQLException throwables) {
                    throwables.printStackTrace();
                }
            }
            if (connection != null) {
                try {
                    connection.close();
                } catch (SQLException throwables) {
                    throwables.printStackTrace();
                }
            }
        }
    }
}
